import numpy as np

class RBF:
    """
    单一RBF核函数
    """
    def __init__(self, center, sigma):
        self.center = np.asarray(center)
        self.sigma = sigma

    def __call__(self, x):
        x = np.atleast_2d(x)
        diff = x - self.center
        dist_sq = np.sum(diff ** 2, axis=1)
        return np.exp(-dist_sq / (2 * self.sigma ** 2))

class RBFN:
    """
    支持分组建模、分组预测和加权组合的RBFN（NumPy实现，适配cc-ddea流程）
    """
    def __init__(self, input_dim, output_dim=1, num_centers=10, sigma=1.0, group_rules=None, random_state=None):
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.num_centers = num_centers
        self.sigma = sigma
        self.random_state = random_state
        self.group_rules = group_rules if group_rules is not None else [[i] for i in range(input_dim)]
        self.groups = self._init_groups(self.group_rules)
        self.base_models = self._init_base_models()
        self.group_weights = np.ones(len(self.groups)) / len(self.groups)  # 初始等权

    def _init_groups(self, group_rules):
        return group_rules

    def _init_base_models(self):
        base_models = []
        for group in self.groups:
            model = _SingleRBFN(len(group), self.output_dim, self.num_centers, self.sigma, self.random_state)
            base_models.append(model)
        return base_models

    def fit(self, X, y):
        for i, group in enumerate(self.groups):
            X_group = X[:, group]
            self.base_models[i].fit(X_group, y)
        errors = np.zeros(len(self.groups))
        for i, group in enumerate(self.groups):
            X_group = X[:, group]
            y_pred = self.base_models[i].predict(X_group)
            errors[i] = np.mean((y - y_pred) ** 2)
        errors = np.clip(errors, 1e-8, None)
        inv_errors = 1.0 / errors
        self.group_weights = inv_errors / np.sum(inv_errors)

    def predict(self, X):
        preds = []
        for i, group in enumerate(self.groups):
            X_group = X[:, group]
            y_pred = self.base_models[i].predict(X_group)
            preds.append(y_pred)
        preds = np.stack(preds, axis=-1)
        weighted = np.tensordot(preds, self.group_weights, axes=([2],[0]))
        return weighted

    def calc_activations(self, X):
        # 返回所有分组拼接的特征 (N, num_centers * n_groups)
        features = []
        for i, group in enumerate(self.groups):
            X_group = X[:, group]
            group_feat = self.base_models[i].calc_activations(X_group)
            features.append(group_feat)
        return np.concatenate(features, axis=1)

class _SingleRBFN:
    """
    单分组RBFN子模型（内部使用）
    """
    def __init__(self, input_dim, output_dim=1, num_centers=10, sigma=1.0, random_state=None):
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.num_centers = num_centers
        self.sigma = sigma
        rng = np.random.default_rng(random_state)
        self.centers = rng.uniform(-5, 5, (num_centers, input_dim))
        self.weights = rng.normal(0, 1, (num_centers, output_dim))
        self.rbfs = [RBF(center, sigma) for center in self.centers]

    def calc_activations(self, X):
        X = np.atleast_2d(X)
        G = np.zeros((X.shape[0], self.num_centers))
        for i, rbf in enumerate(self.rbfs):
            G[:, i] = rbf(X)
        return G

    def fit(self, X, y):
        G = self.calc_activations(X)
        self.weights = np.linalg.pinv(G) @ y

    def predict(self, X):
        G = self.calc_activations(X)
        return G @ self.weights 